# -*- coding: utf-8 -*-
'''
Created on 25.06.2019

@author: yu03
'''

import cv2
import matplotlib.pyplot as plt
import matplotlib
from scipy import signal
from FFT_Interpolation import FFT_interpolation_boxcar, FFT_cal
import numpy as np
from mpl_toolkits.mplot3d import Axes3D


Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

# pattern_path = 'test_far.bmp'
pattern_path = r'F:\Data_Liang_Yu\Users\yu03\Desktop\3-DoF Interferometer\NMC Room\6 Dof PI\Naked_double_Far\Cam\Nonlinearity\weak_ring\Mea.png'
img = cv2.imread(pattern_path, 0)
size = img.shape
print(size)
# cv2.imshow('image',img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
plt.figure('Origin')
im = plt.imshow(img, cmap='jet')
plt.colorbar(im, shrink=0.6)
plt.show()

fig = plt.figure('Spot')
ax = Axes3D(fig)
dx = np.linspace(1, len(img[0]), len(img[0]))
dy = np.linspace(1, len(img), len(img))
X, Y = np.meshgrid(dx, dy)
norm = matplotlib.colors.Normalize(vmin=0, vmax=255)
surf = ax.plot_surface(X, Y, img, rcount=200, ccount=200, cmap='jet', norm=norm)
ax.set_zlim(0, 255)
plt.colorbar(surf, shrink=0.6)
plt.show()

plt.figure('Selected Area')
img = img[200:650,350:700]
plt.imshow(img, cmap='jet')
# plt.show()

hor_magnitude_set = []
ver_magnitude_set = []
hor_magnitude_est_set = []
ver_magnitude_est_set = []
hor_angle_set = []
hor_phase_set = []
ver_angle_set = []
ver_phase_set = []

for i in range(100):
#     hor_center = img[416-50+i][350:700]
#     ver_center = img[:,540-50+i][200:650]
    
    hor_center = img[416-200-50+i]
    ver_center = img[:,540-350-50+i]
        
#     hor_center = hor_center / signal.gaussian(len(hor_center), std=len(hor_center)/2.5)
#     ver_center = ver_center / signal.gaussian(len(ver_center), std=len(ver_center)/2.5)
    
    
    ### Remove DC
    hor_center = hor_center - np.average(hor_center)
    ver_center = ver_center - np.average(ver_center)
    
    ### Coarse FFT
    hor_freqline, hor_fft, hor_magnitude = FFT_cal(hor_center, pix_size)[:3]
    ver_freqline, ver_fft, ver_magnitude = FFT_cal(ver_center, pix_size)[:3]

    hor_magnitude_set.append(hor_magnitude)
    ver_magnitude_set.append(ver_magnitude)
    
    ### FFT Estimation
    hor_freq, hor_phase, hor_freqline_est, hor_magnitude_est = FFT_interpolation_boxcar(hor_center, pix_size)[0:4]
    ver_freq, ver_phase, ver_freqline_est, ver_magnitude_est = FFT_interpolation_boxcar(ver_center, pix_size)[0:4]
    hor_angle = V_x - Lamda * hor_freq / 2
    ver_angle = V_y - Lamda * ver_freq / 2
    print(hor_freq, hor_angle)
    hor_phase = hor_phase / 2 / np.pi * 360 ### in degree
    if hor_phase < 0:
        hor_phase += 360
    ver_phase = ver_phase / 2 / np.pi * 360 ### in degree
    if ver_phase < 0:
        ver_phase += 360
#     print(hor_angle, ver_angle, hor_phase, ver_phase)
    hor_angle_set.append(hor_angle)
    hor_phase_set.append(hor_phase)
    ver_angle_set.append(ver_angle)
    ver_phase_set.append(ver_phase)
    hor_magnitude_est_set.append(hor_magnitude_est)
    ver_magnitude_est_set.append(ver_magnitude_est)
    
plt.figure('Line Shape')
plt.subplot(2,2,1)
plt.plot(hor_center)
plt.subplot(2,2,2)
plt.plot(ver_center)
plt.subplot(2,2,3)
plt.stem(hor_freqline, hor_magnitude, use_line_collection=True)
plt.subplot(2,2,4)
plt.stem(ver_freqline, ver_magnitude, use_line_collection=True)
# plt.show()

hor_magnitude_set = np.array(hor_magnitude_set)
ver_magnitude_set = np.array(ver_magnitude_set)
hor_magnitude_est_set = np.array(hor_magnitude_est_set)
ver_magnitude_est_set = np.array(ver_magnitude_est_set)

fig = plt.figure('Line FFT Compare')
ax1 = fig.add_subplot(2,1,1, projection='3d')
plt.title('Horizontal')
X = np.linspace(hor_freqline[0], hor_freqline[-1], len(hor_freqline))
Y = np.linspace(1, len(hor_magnitude_set), len(hor_magnitude_set))
X, Y = np.meshgrid(X[:30], Y)
surf = ax1.plot_surface(X, Y, hor_magnitude_set[:,:30], rcount=300, ccount=300, cmap='jet')

ax2 = fig.add_subplot(2,1,2, projection='3d')
X = np.linspace(ver_freqline[0], ver_freqline[-1], len(ver_freqline))
Y = np.linspace(1, len(ver_magnitude_set), len(ver_magnitude_set))
X, Y = np.meshgrid(X[:30], Y)
surf = ax2.plot_surface(X, Y, ver_magnitude_set[:,:30], rcount=300, ccount=300, cmap='jet')

# plt.show()
# 
plt.figure('Line Estimation')
plt.subplot(2,1,1)
plt.stem(hor_freqline_est, hor_magnitude_est, use_line_collection=True)
plt.subplot(2,1,2)
plt.stem(ver_freqline_est, ver_magnitude_est, use_line_collection=True)
# plt.show()
 
fig = plt.figure('Line Estimation Compare')
ax1 = fig.add_subplot(2,1,1, projection='3d')
plt.title('Horizontal')
X = np.linspace(hor_freqline_est[0], hor_freqline_est[-1], len(hor_freqline_est))
Y = np.linspace(1, len(hor_magnitude_est_set), len(hor_magnitude_est_set))
X, Y = np.meshgrid(X[:30], Y)
surf = ax1.plot_surface(X, Y, hor_magnitude_est_set[:,:30], rcount=300, ccount=300, cmap='jet')
  
ax2 = fig.add_subplot(2,1,2, projection='3d')
plt.title('Vertical')
X = np.linspace(ver_freqline_est[0], ver_freqline_est[-1], len(ver_freqline_est))
Y = np.linspace(1, len(ver_magnitude_est_set), len(ver_magnitude_est_set))
X, Y = np.meshgrid(X[:30], Y)
surf = ax2.plot_surface(X, Y, ver_magnitude_est_set[:,:30], rcount=300, ccount=300, cmap='jet')
  
plt.show()
  
  
plt.figure('Lines Result Compare')
plt.subplot(2,2,1)
plt.plot(hor_angle_set)
plt.title("Horizontal Tilting")
plt.ylabel("Horizontal Tilting (rad)")
plt.xlabel("Horizontal Lines")
plt.grid(which='major', axis='both')
   
plt.subplot(2,2,3)
plt.plot(hor_phase_set)
plt.title("Horizontal Phase")
plt.ylabel("Horizontal Phase (degree)")
plt.xlabel("Horizontal Lines")
plt.grid(which='major', axis='both')
   
plt.subplot(2,2,2)
plt.plot(ver_angle_set)
plt.title("Vertical Tilting")
plt.ylabel("Vertical Tilting (rad)")
plt.xlabel("Vertical Lines")
plt.grid(which='major', axis='both')
   
plt.subplot(2,2,4)
plt.plot(ver_phase_set)
plt.title("Vertical Phase")
plt.ylabel("Vertical Phase (degree)")
plt.xlabel("Vertical Lines")
plt.grid(which='major', axis='both')
   
plt.tight_layout()

plt.show()